from typing import Any, Dict, Iterator, List, Optional

from langchain_core.callbacks import (
    CallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage
)
from langchain_core.messages.ai import UsageMetadata
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from pydantic import Field


class ChatParrotLink(BaseChatModel):
    """自定义聊天模型，会回显输入消息的前 `parrot_buffer_length` 个字符。

    在向 LangChain 贡献实现时，需详细记录该模型的功能，包括初始化参数，并提供示例代码，
    方便用户理解。此外，还应包含相关文档或 API 的链接。

    示例：

        .. code-block:: python

            model = ChatParrotLink(parrot_buffer_length=2, model="bird-brain-001")
            result = model.invoke([HumanMessage(content="hello")])
            result = model.batch([[HumanMessage(content="hello")],
                                 [HumanMessage(content="world")]])
    """

    model_name: str = Field(alias="model")
    """模型名称"""
    parrot_buffer_length: int
    """从输入消息的最后一条中提取的字符数量。"""
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    timeout: Optional[int] = None
    stop: Optional[List[str]] = None
    max_retries: int = 2

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """重写 `_generate` 方法，以实现聊天模型的逻辑。

        该方法可以调用 API、本地模型，或其他方式来生成对输入消息的响应。

        Args:
            messages:  由多个消息组成的输入提示（prompt）。
            stop:模型应该停止生成的token字符串列表。
                 如果由于停止token而产生停止，则停止token本身
                 应该作为输出的一部分。目前还没有跨模型强制执行，
                 但这是一个很好的实践，因为它使解析模型的下游输出
                 和理解为什么生成停止变得更加容易。
            run_manager: 用于管理 LLM 回调的运行管理器。
        """
        # 
        # 替换这个为实际逻辑，以从列表messages中生成响应信息。
        last_message = messages[-1]
        tokens = last_message.content[: self.parrot_buffer_length]
        ct_input_tokens = sum(len(message.content) for message in messages)
        ct_output_tokens = len(tokens)
        message = AIMessage(
            content=tokens,
            additional_kwargs={},  # 用于向信息添加额外的有效负载
            response_metadata={  # 用于响应元数据
                "time_in_seconds": 3,
            },
            usage_metadata={
                "input_tokens": ct_input_tokens,
                "output_tokens": ct_output_tokens,
                "total_tokens": ct_input_tokens + ct_output_tokens,
            },
        )
        ##

        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    def _stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        """流式输出模型的响应。

        如果模型支持流式生成（streaming），应该实现此方法。
        如果模型不支持流式生成（steaming），则不需要实现它。
        LangChain会自动使用`_generate` 方法处理流式请求。

        Args:
            messages: 由多个消息组成的输入提示（prompt）。
            stop: 模型应该停止生成的token字符串列表。
                  如果由于停止token而产生停止，则停止token本身
                  应该作为输出的一部分。目前还没有跨模型强制执行，
                  但这是一个很好的实践，因为它使解析模型的下游输出
                  和理解为什么生成停止变得更加容易。
            run_manager: 用于管理 LLM 回调的运行管理器。
        """
        last_message = messages[-1]
        tokens = str(last_message.content[: self.parrot_buffer_length])
        ct_input_tokens = sum(len(message.content) for message in messages)

        for token in tokens:
            usage_metadata = UsageMetadata(
                {
                    "input_tokens": ct_input_tokens,
                    "output_tokens": 1,
                    "total_tokens": ct_input_tokens + 1,
                }
            )
            ct_input_tokens = 0
            chunk = ChatGenerationChunk(
                message=AIMessageChunk(content=token, usage_metadata=usage_metadata)
            )

            if run_manager:
                # 这在较新版本的LangChain中是可选的
                # on_llm_new_token将自动调用
                run_manager.on_llm_new_token(token, chunk=chunk)

            yield chunk

        # 让我们添加一些其他信息（例如，响应元数据）
        chunk = ChatGenerationChunk(
            message=AIMessageChunk(content="", response_metadata={"time_in_sec": 3})
        )
        if run_manager:
            # 这在较新版本的LangChain中是可选的
            # on_llm_new_token将自动调用
            run_manager.on_llm_new_token(token, chunk=chunk)
        yield chunk

    @property
    def _llm_type(self) -> str:
        """获取该聊天模型的 LLM 类型。"""
        return "echoing-chat-model-advanced"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """返回包含模型识别参数的字典。

        此信息由LangChain回调系统使用，该系统
        用于跟踪目的，使监视llm成为可能。
        """
        return {
            # 模型名运行用户在LLM监控应用程序中指定
            # 自定义tokens计数规则(e.g., 在LangSmith中，
            # 用户可以为他们的模型提供每个token的定价，
            # 并监控给定LLM的成本）
            "model_name": self.model_name,
        }
if __name__=="__main__":
    model = ChatParrotLink(parrot_buffer_length=20, model="my_custom_model")

    print(model.invoke(
        [
            HumanMessage(content="hello!"),
            AIMessage(content="Hi there human!"),
            HumanMessage(content="Meow!"),
        ]
    ))
    print(model.invoke("hello"))
    print(model.batch(["hello", "goodbye"]))
    for chunk in model.stream("cat"):
        print(chunk.content, end="|")
    import asyncio
    async def ff():
        async for event in model.astream_events("cat", version="v1"):
            print(event)
    asyncio.run(ff())